import cv2
import numpy as np
import sys


# 帮助信息函数
def help():
    print("\nThis is a demo program shows how perspective transformation applied on an image, \n"
          f"Using OpenCV version {cv2.__version__}\n")
    print(f"Usage:\n{sys.argv[0]} [image_name -- Default right.jpg]\n")
    print("\nHot keys: \n"
          "\tESC, q - quit the program\n"
          "\tr - change order of points to rotate transformation\n"
          "\tc - delete selected points\n"
          "\ti - change order of points to inverse transformation \n"
          "\nUse your mouse to select a point and move it to see transformation changes")


# 透视变换函数
def warping(image, warped_image_size, srcPoints, dstPoints):
    M = cv2.getPerspectiveTransform(np.float32(srcPoints), np.float32(dstPoints))
    return cv2.warpPerspective(image, M, warped_image_size)


# 鼠标回调函数
def onMouse(event, x, y, flags, param):
    global roi_corners, dragging, selected_corner_index, validation_needed
    global transformed_points, g_po, g_mx, g_my

    # 当已选择4个点时
    if len(roi_corners) == 4:
        for i in range(4):
            if (event == cv2.EVENT_LBUTTONDOWN and
                    abs(roi_corners[i][0] - x) < 10 and
                    abs(roi_corners[i][1] - y) < 10):
                selected_corner_index = i
                dragging = True

        g_mx, g_my = x, y
        transformed_points = []
        g_po = [
            (original_image_cols / 2.0, original_image_rows),
            (x, y)
        ]
        M = cv2.getPerspectiveTransform(np.float32(roi_corners), np.float32(dst_corners))
        transformed_points = cv2.perspectiveTransform(np.float32([g_po]), M)[0].tolist()
        print(f"{g_mx}:{g_my}")
        print(f"{transformed_points[1][0]}:{transformed_points[1][1]}")
        print(
            f"{transformed_points[1][0] - transformed_points[0][0]}:{-(transformed_points[1][1] - transformed_points[0][1])}")
        print()

    # 添加新点
    elif event == cv2.EVENT_LBUTTONDOWN:
        roi_corners.append((x, y))
        validation_needed = True

    # 释放鼠标
    if event == cv2.EVENT_LBUTTONUP:
        dragging = False

    # 拖动点
    if event == cv2.EVENT_MOUSEMOVE and dragging:
        roi_corners[selected_corner_index] = (x, y)
        validation_needed = True


# 初始化变量
windowTitle = "Perspective Transformation Demo"
labels = ["TL", "TR", "BR", "BL"]
roi_corners = []
midpoints = [(0, 0)] * 4
dst_corners = [(0, 0)] * 4
roiIndex = 0
dragging = False
selected_corner_index = 0
validation_needed = True
transformed_points = []
g_po = []
g_mx, g_my = 0, 0

# 主程序
if __name__ == "__main__":
    help()

    # 解析命令行参数
    filename = "./warpPerspective.jpg"
    if len(sys.argv) > 1:
        filename = sys.argv[1]

    try:
        original_image = cv2.imread(filename)
        if original_image is None:
            print(f"Error: Could not load image {filename}")
            sys.exit(1)
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)

    original_image_cols = float(original_image.shape[1])
    original_image_rows = float(original_image.shape[0])
    print(f"width:{original_image_cols} height:{original_image_rows}")

    # 初始化ROI角点
    roi_corners = [
        (427, 416),
        (662, 416),
        (1072, 665),
        (0, 665)
    ]

    # 创建窗口
    cv2.namedWindow(windowTitle, cv2.WINDOW_NORMAL)
    cv2.namedWindow("Warped Image", cv2.WINDOW_AUTOSIZE)
    cv2.moveWindow("Warped Image", 20, 20)
    cv2.moveWindow(windowTitle, 330, 20)

    # 设置鼠标回调
    cv2.setMouseCallback(windowTitle, onMouse)

    # 主循环
    endProgram = False
    while not endProgram:
        # 当点不足4个时
        if validation_needed and len(roi_corners) < 4:
            validation_needed = False
            image = original_image.copy()

            for i in range(len(roi_corners)):
                cv2.circle(image, roi_corners[i], 5, (0, 255, 0), 3)

                if i > 0:
                    cv2.line(image, roi_corners[i - 1], roi_corners[i], (0, 0, 255), 2)
                    cv2.putText(image, labels[i], roi_corners[i],
                                cv2.FONT_ITALIC, 0.8, (255, 0, 0), 1)

            cv2.imshow(windowTitle, image)

        # 当点等于4个时
        if validation_needed and len(roi_corners) == 4:
            validation_needed = False
            image = original_image.copy()

            # 绘制ROI多边形和角点
            for i in range(4):
                cv2.line(image, roi_corners[i], roi_corners[(i + 1) % 4], (0, 0, 255), 2)
                cv2.circle(image, roi_corners[i], 5, (0, 255, 0), 3)
                cv2.putText(image, labels[i], roi_corners[i],
                            cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 0, 0), 2)

            # 计算中点
            midpoints[0] = ((roi_corners[0][0] + roi_corners[1][0]) / 2,
                            (roi_corners[0][1] + roi_corners[1][1]) / 2)
            midpoints[1] = ((roi_corners[1][0] + roi_corners[2][0]) / 2,
                            (roi_corners[1][1] + roi_corners[2][1]) / 2)
            midpoints[2] = ((roi_corners[2][0] + roi_corners[3][0]) / 2,
                            (roi_corners[2][1] + roi_corners[3][1]) / 2)
            midpoints[3] = ((roi_corners[3][0] + roi_corners[0][0]) / 2,
                            (roi_corners[3][1] + roi_corners[0][1]) / 2)

            # 计算目标角点
            dst_corners[0] = (0, 0)
            dst_corners[1] = (float(np.linalg.norm(np.array(midpoints[1]) - np.array(midpoints[3]))), 0)
            dst_corners[2] = (dst_corners[1][0], float(np.linalg.norm(np.array(midpoints[0]) - np.array(midpoints[2]))))
            dst_corners[3] = (0, dst_corners[2][1])

            # 设置变换后图像尺寸
            warped_image_size = (int(original_image_cols), int(original_image_rows))

            # 计算透视变换矩阵并应用变换
            M = cv2.getPerspectiveTransform(np.float32(roi_corners), np.float32(dst_corners))
            warped_image = cv2.warpPerspective(original_image, M, warped_image_size)

            # 添加文本信息
            text1 = f"TL({int(roi_corners[0][0])}:{int(roi_corners[0][1])})" + \
                    f"TR({int(roi_corners[1][0])}:{int(roi_corners[1][1])})"
            cv2.putText(image, text1, (30, 30),
                        cv2.FONT_ITALIC, 0.8, (255, 0, 0), 1)

            if len(transformed_points) > 1:
                text2 = f"({int(transformed_points[1][0] - transformed_points[0][0])}:" + \
                        f"{int(-(transformed_points[1][1] - transformed_points[0][1]))})"
                cv2.putText(image, text2, (g_mx, g_my),
                            cv2.FONT_ITALIC, 1, (255, 0, 0), 2)

            cv2.imshow(windowTitle, image)
            cv2.imshow("Warped Image", warped_image)

        # 处理键盘事件
        c = cv2.waitKey(10) & 0xFF

        if c in [ord('q'), ord('Q'), 27]:  # ESC, q, Q
            endProgram = True

        if c in [ord('c'), ord('C')]:  # 清除所有点
            roi_corners = []

        if c in [ord('r'), ord('R')]:  # 旋转点顺序
            if len(roi_corners) > 0:
                roi_corners = roi_corners[1:] + [roi_corners[0]]
                validation_needed = True

        if c in [ord('i'), ord('I')]:  # 反转点顺序
            if len(roi_corners) == 4:
                roi_corners[0], roi_corners[1] = roi_corners[1], roi_corners[0]
                roi_corners[2], roi_corners[3] = roi_corners[3], roi_corners[2]
                validation_needed = True

        # 检查窗口是否关闭
        if cv2.getWindowProperty(windowTitle, cv2.WND_PROP_VISIBLE) < 1:
            break
        if cv2.getWindowProperty("Warped Image", cv2.WND_PROP_VISIBLE) < 1:
            break

    cv2.destroyAllWindows()